package org.dreamwork.sso.saml2.protocol;

import org.dreamwork.sso.saml2.util.SecretHelper;
import org.dreamwork.misc.XMLUtil;
import org.dreamwork.misc.Base64;
import org.joda.time.DateTime;
import org.opensaml.Configuration;
import org.opensaml.common.SAMLVersion;
import org.opensaml.saml2.core.*;
import org.opensaml.saml2.core.impl.IssuerBuilder;
import org.opensaml.saml2.core.impl.ResponseBuilder;
import org.opensaml.saml2.core.impl.StatusBuilder;
import org.opensaml.saml2.core.impl.StatusCodeBuilder;
import org.opensaml.saml2.encryption.Decrypter;
import org.opensaml.saml2.encryption.Encrypter;
import org.opensaml.xml.XMLObjectBuilderFactory;
import org.opensaml.xml.encryption.EncryptionConstants;
import org.opensaml.xml.encryption.EncryptionParameters;
import org.opensaml.xml.encryption.InlineEncryptedKeyResolver;
import org.opensaml.xml.encryption.KeyEncryptionParameters;
import org.opensaml.xml.io.Marshaller;
import org.opensaml.xml.io.MarshallingException;
import org.opensaml.xml.io.Unmarshaller;
import org.opensaml.xml.security.SecurityHelper;
import org.opensaml.xml.security.credential.BasicCredential;
import org.opensaml.xml.security.keyinfo.KeyInfoCredentialResolver;
import org.opensaml.xml.security.keyinfo.KeyInfoHelper;
import org.opensaml.xml.security.keyinfo.StaticKeyInfoCredentialResolver;
import org.opensaml.xml.security.keyinfo.StaticKeyInfoGenerator;
import org.opensaml.xml.security.x509.BasicX509Credential;
import org.opensaml.xml.security.x509.X509KeyInfoGeneratorFactory;
import org.opensaml.xml.signature.*;
import org.opensaml.xml.signature.impl.KeyInfoBuilder;
import org.opensaml.xml.signature.impl.KeyNameBuilder;
import org.opensaml.xml.signature.impl.SignatureBuilder;
import org.opensaml.xml.validation.ValidationException;
import org.opensaml.xml.validation.Validator;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import javax.crypto.SecretKey;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;

/**
 * Created with IntelliJ IDEA.
 * User: seth.yang
 * Date: 13-8-13
 * Time: 下午5:10
 */
public class SamlLoginResponse extends SamlObject<Response> {
    private Response response;

    private SamlAssertion assertion;

    public SamlLoginResponse (XMLObjectBuilderFactory factory) {
        super(factory);

        response = new ResponseBuilder ().buildObject ();
        response.setVersion (SAMLVersion.VERSION_20);
        setIssueInstant (new DateTime ());
    }

    public SamlLoginResponse (String encodedContent) {
        super (Configuration.getBuilderFactory ());
        try {
            encodedContent = encodedContent.replace (' ', '+');
            Document doc = XMLUtil.parse (new ByteArrayInputStream (Base64.decode (encodedContent)));
            Element e = doc.getDocumentElement ();
            Unmarshaller unmarshaller = Configuration.getUnmarshallerFactory ().getUnmarshaller (e);
            response = (Response) unmarshaller.unmarshall (e);

            List<Assertion> list = response.getAssertions ();
            if (list != null && list.size () > 0) {
                setAssertion (new SamlAssertion (list.get (0)));
            }
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }
    }

    public SamlLoginResponse (XMLObjectBuilderFactory factory, SamlLoginRequest request) {
        this (factory);
        setId (request.getId ());
//        setIssuer (request.getIssuer ());
        setDestination (request.getAssertionConsumerServiceURL ());
    }

    public String getId () {
        return response.getID ();
    }

    public void setId (String id) {
        response.setID (id);
    }

    public String getInResponseTo () {
        return response.getInResponseTo ();
    }

    public void setInResponseTo (String newInResponseTo) {
        response.setInResponseTo (newInResponseTo);
    }

    public DateTime getIssueInstant () {
        return response.getIssueInstant ();
    }

    public void setIssueInstant (DateTime issueInstant) {
        response.setIssueInstant (issueInstant);
    }

    public String getDestination () {
        return response.getDestination ();
    }

    public void setDestination (String destination) {
        response.setDestination (destination);
    }

    public String getConsent () {
        return response.getConsent ();
    }

    public void setConsent (String newConsent) {
        response.setConsent (newConsent);
    }

    public String getIssuer () {
        Issuer issuer = response.getIssuer ();
        if (issuer == null) return null;

        return issuer.getValue ();
    }

    public void setIssuer (String issuer) {
        Issuer i = new IssuerBuilder ().buildObject ();
        i.setValue (issuer);
        response.setIssuer (i);
    }

    public SamlStatus getStatus () {
        Status status = response.getStatus ();
        return status == null ? null : SamlStatus.parse (status.getStatusCode ().getValue ());
    }

    public void setStatus (SamlStatus status) {
        Status s = response.getStatus ();
        if (s == null) {
            s = new StatusBuilder ().buildObject ();
            response.setStatus (s);
        }
        StatusCode sc = s.getStatusCode ();
        if (sc == null) {
            sc = new StatusCodeBuilder ().buildObject ();
            s.setStatusCode (sc);
        }
        sc.setValue (status.getStatusCode ());
    }

    public List<Validator> getValidators () {
        return response.getValidators ();
    }

    public void registerValidator (Validator validator) {
        response.registerValidator (validator);
    }

    public void deregisterValidator (Validator validator) {
        response.deregisterValidator (validator);
    }

    @Override
    public Element getDom () {
        Element e = response.getDOM ();
        if (e == null) {
            Marshaller m = Configuration.getMarshallerFactory ().getMarshaller (response);
            try {
                e = m.marshall (response);
            } catch (MarshallingException ex) {
                logger.warn (ex.getMessage (), ex);
                throw new RuntimeException (ex);
            }
        }
        return e;
    }

    public boolean isSigned () {
        return response.isSigned ();
    }

    public boolean verifySignature () {
        if (!response.isSigned ()) {
            if (logger.isDebugEnabled ())
                logger.debug ("The response is not signed.");

            return false;
        }

        try {
            Signature signature = response.getSignature ();
            KeyInfo keyInfo = signature.getKeyInfo ();
            if (keyInfo != null) {
                List<X509Certificate> x509certificates = KeyInfoHelper.getCertificates (keyInfo);

                if ((x509certificates != null) && (!x509certificates.isEmpty ())) {
                    X509Certificate x509 = x509certificates.get (0);
                    BasicX509Credential credential = SecurityHelper.getSimpleCredential (x509, null);
                    if (credential == null) {
                        throw new IllegalArgumentException ("No X.509 certificate to verify signature.");
                    }
                    SignatureValidator validator = new SignatureValidator (credential);
                    try {
                        validator.validate (signature);
                        return true;
                    } catch (ValidationException ex) {
                        logger.warn (ex.getMessage (), ex);
                        return false;
                    }
                }
            }
        } catch (RuntimeException ex) {
            throw ex;
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }

        return false;
    }

    public boolean isAssertionEncrypted () {
        if (response == null) return false;

        List<Assertion> as = response.getAssertions ();
        if (as != null && as.size () > 0) {
            return false;
        }

        List<EncryptedAssertion> eas = response.getEncryptedAssertions ();
        return eas != null && eas.size () > 0;

    }

    public SamlAssertion getAssertion () {
        return assertion;
    }

    public void setAssertion (SamlAssertion assertion) {
        this.assertion = assertion;
    }

    public SamlAssertion decryptAssertion (PrivateKey privateKey) {
        try {
            List<EncryptedAssertion> eas = response.getEncryptedAssertions ();
            if (eas == null || eas.size () == 0) return null;

            EncryptedAssertion ea = eas.get (0);
            BasicX509Credential x509 = new BasicX509Credential ();
            x509.setPrivateKey (privateKey);
            KeyInfoCredentialResolver resolver = new StaticKeyInfoCredentialResolver (x509);
            InlineEncryptedKeyResolver kekEncResolver = new InlineEncryptedKeyResolver ();
            Decrypter decrypter = new Decrypter (null, resolver, kekEncResolver);
            Assertion assertion = decrypter.decrypt (ea);
            return this.assertion = new SamlAssertion (assertion);
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }

    public void sign (X509Certificate cert, PrivateKey privateKeyOfCert) {
        if (assertion == null)
            throw new RuntimeException ("Assert is null");

        if (cert == null)
            throw new RuntimeException ("Certificate can't be null");

        if (privateKeyOfCert == null)
            throw new RuntimeException ("private key of certificate can't be null");

        try {
            BasicX509Credential x509 = new BasicX509Credential ();
            x509.setPrivateKey (privateKeyOfCert);
            x509.setPublicKey (cert.getPublicKey ());
            x509.setEntityCertificate (cert);

            SignatureBuilder signatureBuilder = new SignatureBuilder ();
            Signature signature = signatureBuilder.buildObject();
            signature.setSigningCredential(x509);
            signature.setCanonicalizationAlgorithm(SignatureConstants.ALGO_ID_C14N_EXCL_OMIT_COMMENTS);
            signature.setSignatureAlgorithm(SignatureConstants.ALGO_ID_SIGNATURE_RSA);

            X509KeyInfoGeneratorFactory f = new X509KeyInfoGeneratorFactory ();
            f.setEmitEntityCertificate (true);
            KeyInfo keyInfo = f.newInstance ().generate (x509);
            signature.setKeyInfo (keyInfo);
            response.setSignature(signature);
            getDom (); // ???
            Signer.signObject (signature);

//            verifySignature ();
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }

    public void encrypt (PublicKey publicKey) {
        if (assertion == null)
            throw new RuntimeException ("Assertion is null");

        if (publicKey == null)
            throw new RuntimeException ("Public key can't be null");
        try {

            KeyInfoBuilder builder = new KeyInfoBuilder ();

            // generate random aes key
            SecretKey aesKey = SecretHelper.createAESKey ();
            BasicCredential aes = new BasicCredential ();
            aes.setSecretKey (aesKey);

            // prepare to encrypting THE aes key
            BasicCredential x509 = new BasicCredential ();
            x509.setPublicKey (publicKey);

            // prepare to encrypting content
            EncryptionParameters params = new EncryptionParameters();
            params.setAlgorithm (EncryptionConstants.ALGO_ID_BLOCKCIPHER_AES128);
            params.setEncryptionCredential (aes);

            KeyEncryptionParameters kekParamsRSA = new KeyEncryptionParameters();
            kekParamsRSA.setAlgorithm(EncryptionConstants.ALGO_ID_KEYTRANSPORT_RSAOAEP);
            kekParamsRSA.setEncryptionCredential(x509);

            List<KeyEncryptionParameters> kekParamsList = new ArrayList<KeyEncryptionParameters> ();

            KeyInfo kekKeyInfoRSA = builder.buildObject ();

            KeyName keyName = new KeyNameBuilder ().buildObject ();
            keyName.setValue(assertion.getIssuer ());
            kekKeyInfoRSA.getKeyNames().add(keyName);
            kekParamsRSA.setKeyInfoGenerator(new StaticKeyInfoGenerator(kekKeyInfoRSA));
            kekParamsList.add(kekParamsRSA);

            Encrypter encrypter = new Encrypter (params, kekParamsList);
            encrypter.setKeyPlacement(Encrypter.KeyPlacement.INLINE);
            EncryptedAssertion ea = encrypter.encrypt (assertion.getAssertion ());

            response.getEncryptedAssertions ().add (ea);
        } catch (Exception ex) {
            logger.warn (ex.getMessage (), ex);
            throw new RuntimeException (ex);
        }
    }

    public byte[] getSignatureKeyInfo () {
        KeyInfo ki = response.getSignature ().getKeyInfo ();
        Element e = ki.getDOM ();
        if (e == null) {
            Marshaller m = Configuration.getMarshallerFactory ().getMarshaller (ki);
            try {
                e = m.marshall (ki);
            } catch (MarshallingException ex) {
                logger.warn (ex.getMessage (), ex);
                throw new RuntimeException (ex);
            }
        }
        try {
            ByteArrayOutputStream baos = new ByteArrayOutputStream ();
            XMLUtil.print (e, baos);
            return baos.toByteArray ();
        } catch (Exception ex) {
            throw new RuntimeException (ex);
        }
    }
}